"""
AI服务路由器
负责将视频帧路由到相应的AI服务，实现负载均衡和故障转移
"""

import asyncio
import logging
import time
from typing import Dict, List, Any, Optional, Callable, Set
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict, deque
import numpy as np
from datetime import datetime, timedelta
import uuid

logger = logging.getLogger(__name__)


class ServiceStatus(Enum):
    """服务状态"""
    AVAILABLE = "available"
    BUSY = "busy"
    ERROR = "error"
    MAINTENANCE = "maintenance"


@dataclass
class AIServiceInfo:
    """AI服务信息"""
    service_name: str
    service_instance: Any
    status: ServiceStatus = ServiceStatus.AVAILABLE
    current_load: int = 0
    max_concurrent: int = 10
    avg_processing_time: float = 0.0
    success_rate: float = 1.0
    last_health_check: datetime = field(default_factory=datetime.now)
    error_count: int = 0
    total_requests: int = 0
    
    @property
    def load_percentage(self) -> float:
        """负载百分比"""
        return (self.current_load / max(self.max_concurrent, 1)) * 100
    
    @property
    def is_available(self) -> bool:
        """是否可用"""
        return (self.status == ServiceStatus.AVAILABLE and 
                self.current_load < self.max_concurrent)


@dataclass
class RoutingRequest:
    """路由请求"""
    request_id: str
    camera_id: str
    algorithm: str
    frame: np.ndarray
    priority: int = 5  # 1-10, 10最高
    timeout: float = 30.0
    callback: Optional[Callable] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)


@dataclass
class RoutingResult:
    """路由结果"""
    request_id: str
    success: bool
    result: Any = None
    error_message: Optional[str] = None
    service_name: Optional[str] = None
    processing_time: float = 0.0
    completed_at: datetime = field(default_factory=datetime.now)


class LoadBalancingStrategy(Enum):
    """负载均衡策略"""
    ROUND_ROBIN = "round_robin"
    LEAST_CONNECTIONS = "least_connections"
    WEIGHTED_ROUND_ROBIN = "weighted_round_robin"
    RESPONSE_TIME = "response_time"
    RANDOM = "random"


class AIServiceRouter:
    """AI服务路由器"""
    
    def __init__(self, config: Dict[str, Any] = None):
        self.config = config or {}
        
        # 服务注册表
        self.services: Dict[str, List[AIServiceInfo]] = defaultdict(list)
        
        # 路由策略
        self.load_balancing_strategy = LoadBalancingStrategy(
            self.config.get("load_balancing_strategy", "least_connections")
        )
        
        # 故障转移配置
        self.max_retries = self.config.get("max_retries", 3)
        self.retry_delay = self.config.get("retry_delay", 1.0)
        self.circuit_breaker_threshold = self.config.get("circuit_breaker_threshold", 0.5)
        
        # 健康检查配置
        self.health_check_interval = self.config.get("health_check_interval", 60)
        self.health_check_timeout = self.config.get("health_check_timeout", 10)
        
        # 摄像头配置缓存
        self.camera_configs: Dict[str, Dict[str, Any]] = {}
        
        # 统计信息
        self.stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "avg_response_time": 0.0,
            "service_failures": defaultdict(int),
            "load_balancing_decisions": defaultdict(int)
        }
        
        # 运行状态
        self.is_running = False
        self.health_check_task = None
        
        # 轮询计数器
        self.round_robin_counters = defaultdict(int)
        
        # 熔断器状态
        self.circuit_breakers: Dict[str, Dict[str, Any]] = defaultdict(dict)
        
        logger.info("AI服务路由器已初始化")
    
    def register_service(self, algorithm: str, service_instance: Any, 
                        max_concurrent: int = 10, service_name: str = None):
        """注册AI服务"""
        if service_name is None:
            service_name = f"{algorithm}_service_{len(self.services[algorithm])}"
        
        service_info = AIServiceInfo(
            service_name=service_name,
            service_instance=service_instance,
            max_concurrent=max_concurrent
        )
        
        self.services[algorithm].append(service_info)
        
        logger.info(f"已注册AI服务: {algorithm} -> {service_name}")
    
    def unregister_service(self, algorithm: str, service_name: str):
        """注销AI服务"""
        if algorithm in self.services:
            self.services[algorithm] = [
                s for s in self.services[algorithm] 
                if s.service_name != service_name
            ]
            logger.info(f"已注销AI服务: {algorithm} -> {service_name}")
    
    def set_camera_config(self, camera_id: str, ai_algorithms: List[str], 
                         priority: int = 5, custom_config: Dict[str, Any] = None):
        """设置摄像头AI算法配置"""
        self.camera_configs[camera_id] = {
            "ai_algorithms": ai_algorithms,
            "priority": priority,
            "custom_config": custom_config or {},
            "updated_at": datetime.now()
        }
        
        logger.debug(f"已设置摄像头{camera_id}配置: {ai_algorithms}")
    
    def get_camera_config(self, camera_id: str) -> Dict[str, Any]:
        """获取摄像头配置"""
        return self.camera_configs.get(camera_id, {
            "ai_algorithms": [],
            "priority": 5,
            "custom_config": {},
            "updated_at": datetime.now()
        })
    
    async def route_frame_to_services(self, camera_id: str, frame: np.ndarray, 
                                    frame_timestamp: datetime = None) -> List[RoutingResult]:
        """将视频帧路由到相应的AI服务"""
        if frame_timestamp is None:
            frame_timestamp = datetime.now()
        
        # 获取摄像头配置
        camera_config = self.get_camera_config(camera_id)
        ai_algorithms = camera_config.get("ai_algorithms", [])
        priority = camera_config.get("priority", 5)
        
        if not ai_algorithms:
            logger.warning(f"摄像头{camera_id}未配置AI算法")
            return []
        
        # 创建路由请求
        requests = []
        for algorithm in ai_algorithms:
            request_id = str(uuid.uuid4())
            routing_request = RoutingRequest(
                request_id=request_id,
                camera_id=camera_id,
                algorithm=algorithm,
                frame=frame.copy(),
                priority=priority,
                metadata={
                    "frame_timestamp": frame_timestamp.isoformat(),
                    "camera_config": camera_config.get("custom_config", {})
                }
            )
            requests.append(routing_request)
        
        # 并行处理所有算法
        tasks = [self._process_routing_request(req) for req in requests]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        routing_results = []
        for result in results:
            if isinstance(result, Exception):
                logger.error(f"路由处理异常: {result}")
                continue
            routing_results.append(result)
        
        return routing_results
    
    async def _process_routing_request(self, request: RoutingRequest) -> RoutingResult:
        """处理单个路由请求"""
        start_time = time.time()
        
        try:
            # 选择服务实例
            service_info = await self._select_service(request.algorithm, request.priority)
            if not service_info:
                raise RuntimeError(f"没有可用的{request.algorithm}服务")
            
            # 增加负载计数
            service_info.current_load += 1
            
            try:
                # 调用AI服务
                result = await self._call_ai_service(service_info, request)
                
                # 更新统计
                processing_time = time.time() - start_time
                self._update_service_stats(service_info, processing_time, True)
                self._update_global_stats(processing_time, True)
                
                return RoutingResult(
                    request_id=request.request_id,
                    success=True,
                    result=result,
                    service_name=service_info.service_name,
                    processing_time=processing_time
                )
                
            finally:
                # 减少负载计数
                service_info.current_load = max(0, service_info.current_load - 1)
        
        except Exception as e:
            processing_time = time.time() - start_time
            
            # 更新统计
            if 'service_info' in locals():
                self._update_service_stats(service_info, processing_time, False)
            self._update_global_stats(processing_time, False)
            
            # 尝试故障转移
            if self.max_retries > 0:
                retry_result = await self._retry_with_failover(request, str(e))
                if retry_result:
                    return retry_result
            
            return RoutingResult(
                request_id=request.request_id,
                success=False,
                error_message=str(e),
                processing_time=processing_time
            )
    
    async def _select_service(self, algorithm: str, priority: int = 5) -> Optional[AIServiceInfo]:
        """选择服务实例"""
        if algorithm not in self.services:
            return None
        
        available_services = [
            s for s in self.services[algorithm] 
            if s.is_available and not self._is_circuit_breaker_open(s.service_name)
        ]
        
        if not available_services:
            return None
        
        # 根据负载均衡策略选择服务
        if self.load_balancing_strategy == LoadBalancingStrategy.ROUND_ROBIN:
            return self._select_round_robin(algorithm, available_services)
        
        elif self.load_balancing_strategy == LoadBalancingStrategy.LEAST_CONNECTIONS:
            return min(available_services, key=lambda s: s.current_load)
        
        elif self.load_balancing_strategy == LoadBalancingStrategy.RESPONSE_TIME:
            return min(available_services, key=lambda s: s.avg_processing_time)
        
        elif self.load_balancing_strategy == LoadBalancingStrategy.WEIGHTED_ROUND_ROBIN:
            return self._select_weighted_round_robin(algorithm, available_services, priority)
        
        else:  # RANDOM
            import random
            return random.choice(available_services)
    
    def _select_round_robin(self, algorithm: str, services: List[AIServiceInfo]) -> AIServiceInfo:
        """轮询选择"""
        counter = self.round_robin_counters[algorithm]
        selected = services[counter % len(services)]
        self.round_robin_counters[algorithm] = (counter + 1) % len(services)
        return selected
    
    def _select_weighted_round_robin(self, algorithm: str, services: List[AIServiceInfo], 
                                   priority: int) -> AIServiceInfo:
        """加权轮询选择"""
        # 根据优先级和服务性能计算权重
        weights = []
        for service in services:
            weight = priority * (1 - service.load_percentage / 100) * service.success_rate
            weights.append(max(0.1, weight))  # 最小权重0.1
        
        # 加权随机选择
        import random
        total_weight = sum(weights)
        r = random.uniform(0, total_weight)
        
        cumulative_weight = 0
        for i, weight in enumerate(weights):
            cumulative_weight += weight
            if r <= cumulative_weight:
                return services[i]
        
        return services[-1]  # 备选
    
    async def _call_ai_service(self, service_info: AIServiceInfo, 
                              request: RoutingRequest) -> Any:
        """调用AI服务"""
        service = service_info.service_instance
        
        # 根据算法类型调用相应方法
        if request.algorithm == "behavior_analysis":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'analyze_behavior'):
                return await service.analyze_behavior(request.frame)
        
        elif request.algorithm == "crowd_density":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'analyze_crowd_density'):
                return await service.analyze_crowd_density(request.frame)
        
        elif request.algorithm == "waste_detection":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'detect_waste'):
                return await service.detect_waste(request.frame)
        
        elif request.algorithm == "cleanliness_assessment":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'assess_cleanliness'):
                return await service.assess_cleanliness(request.frame)
        
        elif request.algorithm == "teaching_quality":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'analyze_teaching_quality'):
                return await service.analyze_teaching_quality(request.frame)
        
        elif request.algorithm == "classroom_quality":
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'analyze_classroom_quality'):
                return await service.analyze_classroom_quality(request.frame)
        
        else:
            # 通用调用方式
            if hasattr(service, 'process_video_frame'):
                return await service.process_video_frame(request.frame, request.metadata)
            elif hasattr(service, 'process'):
                return await service.process(request.frame)
            else:
                raise RuntimeError(f"服务{service_info.service_name}不支持算法{request.algorithm}")
    
    async def _retry_with_failover(self, request: RoutingRequest, 
                                  original_error: str) -> Optional[RoutingResult]:
        """故障转移重试"""
        for retry_count in range(self.max_retries):
            try:
                await asyncio.sleep(self.retry_delay * (retry_count + 1))
                
                # 选择其他服务实例
                service_info = await self._select_service(request.algorithm, request.priority)
                if not service_info:
                    continue
                
                # 重试调用
                service_info.current_load += 1
                try:
                    result = await self._call_ai_service(service_info, request)
                    
                    logger.info(f"故障转移成功: {request.algorithm} -> {service_info.service_name}")
                    
                    return RoutingResult(
                        request_id=request.request_id,
                        success=True,
                        result=result,
                        service_name=service_info.service_name,
                        processing_time=0.0  # 重试时间不计入统计
                    )
                    
                finally:
                    service_info.current_load = max(0, service_info.current_load - 1)
                    
            except Exception as e:
                logger.warning(f"故障转移重试{retry_count + 1}失败: {e}")
                continue
        
        return None
    
    def _update_service_stats(self, service_info: AIServiceInfo, 
                             processing_time: float, success: bool):
        """更新服务统计"""
        service_info.total_requests += 1
        
        if success:
            # 更新平均处理时间
            if service_info.avg_processing_time == 0:
                service_info.avg_processing_time = processing_time
            else:
                service_info.avg_processing_time = (
                    service_info.avg_processing_time * 0.9 + processing_time * 0.1
                )
        else:
            service_info.error_count += 1
        
        # 更新成功率
        service_info.success_rate = (
            (service_info.total_requests - service_info.error_count) / 
            max(service_info.total_requests, 1)
        )
        
        # 检查熔断器
        self._check_circuit_breaker(service_info)
    
    def _update_global_stats(self, processing_time: float, success: bool):
        """更新全局统计"""
        self.stats["total_requests"] += 1
        
        if success:
            self.stats["successful_requests"] += 1
            
            # 更新平均响应时间
            if self.stats["avg_response_time"] == 0:
                self.stats["avg_response_time"] = processing_time
            else:
                self.stats["avg_response_time"] = (
                    self.stats["avg_response_time"] * 0.9 + processing_time * 0.1
                )
        else:
            self.stats["failed_requests"] += 1
    
    def _check_circuit_breaker(self, service_info: AIServiceInfo):
        """检查熔断器"""
        service_name = service_info.service_name
        
        # 如果错误率超过阈值，打开熔断器
        if (service_info.total_requests >= 10 and 
            service_info.success_rate < self.circuit_breaker_threshold):
            
            self.circuit_breakers[service_name] = {
                "open": True,
                "opened_at": datetime.now(),
                "failure_count": service_info.error_count
            }
            
            logger.warning(f"熔断器已打开: {service_name}, 成功率: {service_info.success_rate:.2f}")
    
    def _is_circuit_breaker_open(self, service_name: str) -> bool:
        """检查熔断器是否打开"""
        if service_name not in self.circuit_breakers:
            return False
        
        breaker = self.circuit_breakers[service_name]
        if not breaker.get("open", False):
            return False
        
        # 检查是否应该尝试半开状态
        opened_at = breaker.get("opened_at", datetime.now())
        if datetime.now() - opened_at > timedelta(minutes=5):
            # 尝试半开状态
            breaker["open"] = False
            logger.info(f"熔断器尝试半开: {service_name}")
            return False
        
        return True
    
    async def start(self):
        """启动路由器"""
        if self.is_running:
            return
        
        self.is_running = True
        
        # 启动健康检查任务
        self.health_check_task = asyncio.create_task(self._health_check_loop())
        
        logger.info("AI服务路由器已启动")
    
    async def stop(self):
        """停止路由器"""
        self.is_running = False
        
        if self.health_check_task:
            self.health_check_task.cancel()
            try:
                await self.health_check_task
            except asyncio.CancelledError:
                pass
        
        logger.info("AI服务路由器已停止")
    
    async def _health_check_loop(self):
        """健康检查循环"""
        while self.is_running:
            try:
                await self._perform_health_checks()
                await asyncio.sleep(self.health_check_interval)
            except Exception as e:
                logger.error(f"健康检查异常: {e}")
                await asyncio.sleep(60)
    
    async def _perform_health_checks(self):
        """执行健康检查"""
        for algorithm, service_list in self.services.items():
            for service_info in service_list:
                try:
                    # 简单的健康检查
                    if hasattr(service_info.service_instance, 'health_check'):
                        is_healthy = await asyncio.wait_for(
                            service_info.service_instance.health_check(),
                            timeout=self.health_check_timeout
                        )
                        
                        if is_healthy:
                            if service_info.status == ServiceStatus.ERROR:
                                service_info.status = ServiceStatus.AVAILABLE
                                logger.info(f"服务恢复: {service_info.service_name}")
                        else:
                            service_info.status = ServiceStatus.ERROR
                            logger.warning(f"服务健康检查失败: {service_info.service_name}")
                    
                    service_info.last_health_check = datetime.now()
                    
                except Exception as e:
                    service_info.status = ServiceStatus.ERROR
                    logger.error(f"健康检查异常 {service_info.service_name}: {e}")
    
    def get_service_stats(self) -> Dict[str, Any]:
        """获取服务统计"""
        service_stats = {}
        
        for algorithm, service_list in self.services.items():
            service_stats[algorithm] = []
            for service_info in service_list:
                stats = {
                    "service_name": service_info.service_name,
                    "status": service_info.status.value,
                    "current_load": service_info.current_load,
                    "load_percentage": service_info.load_percentage,
                    "avg_processing_time": service_info.avg_processing_time,
                    "success_rate": service_info.success_rate,
                    "total_requests": service_info.total_requests,
                    "error_count": service_info.error_count,
                    "last_health_check": service_info.last_health_check.isoformat()
                }
                service_stats[algorithm].append(stats)
        
        return {
            "global_stats": self.stats,
            "service_stats": service_stats,
            "circuit_breakers": {
                name: {
                    "open": breaker.get("open", False),
                    "opened_at": breaker.get("opened_at", datetime.now()).isoformat(),
                    "failure_count": breaker.get("failure_count", 0)
                }
                for name, breaker in self.circuit_breakers.items()
            },
            "camera_configs": {
                camera_id: {
                    **config,
                    "updated_at": config["updated_at"].isoformat()
                }
                for camera_id, config in self.camera_configs.items()
            }
        }


# 全局AI服务路由器实例
ai_service_router = None


def initialize_ai_service_router(config: Dict[str, Any] = None):
    """初始化AI服务路由器"""
    global ai_service_router
    ai_service_router = AIServiceRouter(config)
    logger.info("AI服务路由器已初始化")


def get_ai_service_router() -> Optional[AIServiceRouter]:
    """获取AI服务路由器实例"""
    return ai_service_router


async def shutdown_ai_service_router():
    """关闭AI服务路由器"""
    global ai_service_router
    if ai_service_router:
        await ai_service_router.stop()
        ai_service_router = None
        logger.info("AI服务路由器已关闭")